#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import numpy as np
import json

############################################################################

file_dir = r'/root/data/Erised/ricemee-age/Images'
file_list_file = r'/root/data/Erised/ricemee-age/Annotations/bbox_list.txt'
split_list_file = r'/root/data/Erised/ricemee-age/Annotations/splitlist.txt'
gender_anno_dir = r'/root/data/Erised/ricemee-age/Annotations/gender'

# file_dir = r'/root/data/Erised/age-gender/Images'
# file_list_file = r'/root/data/Erised/age-gender/Annotations/bbox_list.txt'
# split_list_file = r'/root/data/Erised/age-gender/Annotations/splitlist.txt'
# gender_anno_dir = None

SCORE_THR = 0.8


def parse_gender_annotations(gender_anno_file):
    gender_anno_file = gender_anno_file.encode('utf-8')
    gender_label = None
    if os.path.exists(gender_anno_file):
        with open(gender_anno_file, 'r', encoding='utf-8-sig') as json_file:
            label = json.load(json_file)["outputs"]
            if "object" in label:
                label = label["object"][0]['name']
                if label == '女':
                    gender_label = 0
                elif label == '男':
                    gender_label = 1
                elif label == '跳过':
                    pass
                else:
                    raise Exception("Unknown gender label: {}".format(label))

    return gender_label


train_image_set = set()
val_image_set = set()
with open(split_list_file, 'r', encoding='utf-8') as file:
    for line in file.readlines():
        # 33-F/1-2-33.jpg 0
        lines = line.strip().rsplit(' ', 1)
        file_name = lines[0]
        is_val = lines[1] != '0'
        if is_val:
            val_image_set.add(file_name)
        else:
            train_image_set.add(file_name)

file_count = 0
out_path = os.path.dirname(file_list_file)
out_file_path = os.path.join(out_path, 'person_bbox_gt.txt')
out_file = open(out_file_path, 'w', encoding='utf-8')
train_out_file_path = os.path.join(out_path, 'train_person_bbox_gender.txt')
train_out_file = open(train_out_file_path, 'w', encoding='utf-8')
val_out_file_path = os.path.join(out_path, 'val_person_bbox_gender.txt')
val_out_file = open(val_out_file_path, 'w', encoding='utf-8')
with open(file_list_file, 'r', encoding='utf-8') as file:
    for line in file.readlines():
        # 28岁/Screenshot_2018_1127_172021.png 12,33,873,968,0.9995991587638855;93,374,227,746,0.9975289702415466
        file_count += 1
        if file_count % 100 == 0:
            print('Process File Count:', file_count)
        # lines = line.strip().split(' ')
        # count = len(lines)
        # if count > 2:
        #     file_name = ' '.join(lines[:-1])
        # else:
        #     file_name = lines[0]
        lines = line.strip().rsplit(' ', 1)
        file_name = lines[0]
        bbox_info = lines[-1]
        bbox_infos = bbox_info.split(';')
        bbox_list = []
        for bbox_info in bbox_infos:
            bbox = list(map(float, bbox_info.split(',')))
            bbox_list.append(bbox)
        bboxes = np.array(bbox_list, dtype=np.float32)
        max_ares_index = 0
        bbox_count = len(bbox_list)
        if bbox_count > 1:
            mask = bboxes[:, 4] >= SCORE_THR
            w = bboxes[:, 2] - bboxes[:, 0]
            h = bboxes[:, 3] - bboxes[:, 1]
            areas = w * h
            if mask.any():
                bboxes = bboxes[mask]
            if bboxes.shape[0] > 1:
                w = bboxes[:, 2] - bboxes[:, 0]
                h = bboxes[:, 3] - bboxes[:, 1]
                areas = w * h
                max_ares_index = areas.argmax()
        elif bbox_count < 1:
            continue
        bbox = bboxes[max_ares_index]
        need_out = False

        if gender_anno_dir is not None:
            gender_anno_file = os.path.join(gender_anno_dir, os.path.splitext(file_name)[0] + ".json")
            gender = parse_gender_annotations(gender_anno_file)
            if gender is None:
                continue
        else:
            # Extract attributes from image name following 'id-gender-age.jpg' format
            _id, _gender, _age = file_name.split('/')[1].split('.')[0].split('-')
            gender = 1 if int(_gender) == 1 else 0

        out_info = '{} {:.01f},{:.01f},{:.01f},{:.01f} {}\n'.format(file_name, bbox[0], bbox[1], bbox[2], bbox[3], gender)
        if file_name in train_image_set:
            need_out = True
            train_out_file.write(out_info)
        elif file_name in val_image_set:
            need_out = True
            val_out_file.write(out_info)
        if need_out:
            out_file.write('{}\n1\n'.format(file_name))
            out_file.write('{:.01f} {:.01f} {:.01f} {:.01f}\n'.format(bbox[0], bbox[1], bbox[2], bbox[3]))

out_file.close()
train_out_file.close()
val_out_file.close()
os.system('chmod a+wr {}'.format(out_file_path))
os.system('chmod a+wr {}'.format(train_out_file_path))
os.system('chmod a+wr {}'.format(val_out_file_path))
print('Process File Count:', file_count)

print('Finish!')